// File: jmcomic-core/src/main/java/dev/jukomu/common/base/concurrent/QueueProducerConsumer.java
package dev.jukomu.common.base.concurrent;

import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;

/**
 * 对应Python的QueueGenerator，实现生产者-消费者模式，支持停止标志和元素过滤。
 * 同时继承StopFlagManager，以便管理停止信号。
 *
 * @param <E> 队列中元素的类型
 */
public class QueueProducerConsumer<E> extends StopFlagManager implements Iterable<E> {

    // 使用BlockingQueue实现线程安全的队列
    private final LinkedBlockingQueue<E> queue;
    // 用于跟踪所有放入的元素，防止重复
    private final Set<E> totalElements;
    // 用于跟踪已被消费者取出的元素，可以用于某些业务逻辑判断
    private final Set<E> consumedElements;

    private final E endSentinel; // 结束哨兵元素
    private volatile boolean isEnd = false; // 生产者是否已结束生产
    private volatile boolean isGiveback = false; // 消费者是否要求归还当前元素

    private final Predicate<E> filterMethod; // 元素过滤方法

    /**
     * 构造函数。
     *
     * @param endSentinel  标记队列结束的哨兵元素
     * @param stopFlag     用于控制停止的StopFlag实例
     * @param filterMethod 过滤方法，如果为null则不过滤
     */
    public QueueProducerConsumer(E endSentinel, StopFlag stopFlag, Predicate<E> filterMethod) {
        super(stopFlag);
        this.endSentinel = Objects.requireNonNull(endSentinel, "End sentinel cannot be null.");
        this.filterMethod = filterMethod;

        this.queue = new LinkedBlockingQueue<>(); // 无界队列
        this.totalElements = ConcurrentHashMap.newKeySet(); // 线程安全的Set
        this.consumedElements = ConcurrentHashMap.newKeySet(); // 线程安全的Set

        getStopFlag().markRun(); // 标记当前线程为运行状态
    }

    /**
     * 获取所有放入队列的元素总数。
     *
     * @return 元素总数
     */
    public int getTotalCount() {
        return totalElements.size();
    }

    /**
     * 获取队列中当前等待被消费的元素数量。
     *
     * @return 队列大小
     */
    public int getQueueSize() {
        return queue.size();
    }

    /**
     * 获取已消费的元素数量。
     *
     * @return 已消费元素数量
     */
    public int getConsumedCount() {
        return consumedElements.size();
    }

    /**
     * 获取剩余未消费的元素数量。
     *
     * @return 剩余数量
     */
    public int getRemainingCount() {
        return queue.size();
    }

    /**
     * 生产者方法：将元素放入队列。
     *
     * @param obj 要放入的元素
     */
    public void put(E obj) {
        if (Objects.equals(endSentinel, obj)) {
            isEnd = true;
            // 即使是结束哨兵，也放入队列，以便消费者能收到停止信号
            // 确保哨兵只放入一次，防止死循环
            if (!totalElements.contains(endSentinel)) {
                totalElements.add(endSentinel);
                queue.offer(obj); // offer是非阻塞的，如果队列满则返回false，但LinkedBlockingQueue无界
            }
            return;
        }

        if (totalElements.add(obj)) { // 如果元素是第一次被添加到totalElements
            if (filterMethod == null || filterMethod.test(obj)) {
                queue.offer(obj);
            }
        }
    }

    /**
     * 消费者方法：从队列中获取元素，并检查停止标志。
     *
     * @return 获取到的元素
     * @throws InterruptedException 如果线程在等待期间被中断
     */
    private E getAndCheckStop() throws InterruptedException {
        while (true) {
            // 每次尝试获取前都检查停止标志
            if (isStop()) {
                return endSentinel;
            }

            E element = queue.poll(500, TimeUnit.MILLISECONDS); // 尝试获取元素，等待500毫秒

            if (element != null) {
                return element;
            } else {
                // 队列为空，且生产者已结束生产，且队列中没有结束哨兵（可能已取出）
                if (isEnd && queue.isEmpty()) {
                    return endSentinel;
                }
                // 队列为空但生产者未结束，或者队列中有哨兵但尚未取出，继续等待
                // StopFlagManager.sleepOrReturn(0.5) 内部也包含isStop()检查
                sleepOrReturn(0.5);
            }
        }
    }

    /**
     * 请求将当前迭代的元素归还给队列，以便下次迭代再次消费。
     * 仅影响当前迭代，下次迭代会再次返回此元素。
     */
    public void giveback() {
        isGiveback = true;
    }

    /**
     * 通知生产者已完成所有生产，将发送结束哨兵。
     */
    public void close() {
        put(endSentinel);
    }

    @Override
    public Iterator<E> iterator() {
        getStopFlag().markRun(); // 标记当前线程为运行状态
        return new Iterator<E>() {
            private E currentElement = null;

            @Override
            public boolean hasNext() {
                // 如果当前元素需要归还，则hasNext为true
                if (isGiveback) {
                    isGiveback = false; // 重置归还标志
                    return true;
                }
                // 如果已获取到非哨兵元素，则hasNext为true
                if (currentElement != null && !Objects.equals(currentElement, endSentinel)) {
                    return true;
                }
                // 尝试从队列获取下一个元素
                try {
                    currentElement = getAndCheckStop();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt(); // Restore interrupt status
                    currentElement = endSentinel; // Treat as end of stream on interruption
                }

                // 如果获取到哨兵，表示迭代结束
                if (Objects.equals(currentElement, endSentinel)) {
                    currentElement = null; // 清除哨兵，防止后续误用
                    return false;
                }
                // 获取到有效元素
                return true;
            }

            @Override
            public E next() {
                if (!hasNext()) {
                    throw new NoSuchElementException();
                }
                E elementToReturn = currentElement;
                consumedElements.add(elementToReturn); // 标记为已消费
                currentElement = null; // 清除当前元素，以便hasNext获取下一个
                return elementToReturn;
            }
        };
    }

    @Override
    public String toString() {
        return "QueueProducerConsumer{" +
                "totalElements=" + totalElements.size() +
                ", consumedElements=" + consumedElements.size() +
                ", queueSize=" + queue.size() +
                ", isEnd=" + isEnd +
                ", filterMethod=" + (filterMethod != null) +
                '}';
    }
}